import argparse
import tensorflow as tf
from tensorflow import keras as K
import numpy as np
from sklearn.cluster import KMeans


parser = argparse.ArgumentParser()
parser.add_argument("--n_centre", type=int, default=10)
parser.add_argument("--n_class", type=int, default=10)
parser.add_argument("--n_per", type=int, default=512)
args = parser.parse_args()


def euclidean(A, B=None, sqrt=False):
    aTb = np.dot(A, B.T)
    if (B is None) or (B is A):
        aTa = np.diag(aTb)
        bTb = aTa
    else:
        aTa = np.diag(np.dot(A, A.T))
        bTb = np.diag(np.dot(B, B.T))
    D = aTa[:, np.newaxis] - 2.0 * aTb + bTb[np.newaxis, :]
    if sqrt:
        D = np.sqrt(D)
    return D


# data
mnist = K.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0   # [n, 28, 28], [n]
x_train = tf.reshape(x_train, [-1, 784])
#x_test = tf.reshape(x_test, [-1, 784])
print(type(x_train))
x_train, y_train = np.array(x_train), np.array(y_train)
print("data:", type(x_train), x_train.shape, y_test.shape)


# sampling
X_sub, Y_sub = [], []
for c in range(args.n_class):
    Xc = x_train[y_train == c]
    Yc = y_train[y_train == c]
    idx = np.random.permutation(Xc.shape[0])
    Xc, Yc = Xc[idx[:args.n_per]], Yc[idx[:args.n_per]]
    #print("class:", c, ", x:", Xc.shape, "y:", Yc.shape)
    X_sub.append(Xc)
    Y_sub.append(Yc)

X_sub, Y_sub = np.vstack(Xc), np.stack(Yc)
print("X_sub:", X_sub.shape, ", Y_sub:", Y_sub.shape)

# clustering
kmeans = KMeans(n_clusters=args.n_centre, random_state=0).fit(X_sub)
centres = kmeans.cluster_centers_
cluster_id = kmeans.labels_
print("centres:", centres.shape)

variances = []
for i in range(args.n_centre):
    # Var(X) = E(X - EX)^2
    X_id = X_sub[Y_sub == i]
    print("X_id:", i, 
    variances.append(np.mean(euclidean(X_id, centres[i:i+1])))

variances = np.asarray(variances)


# saving
np.save("centres.{}c.{}pc.npy".format(args.n_centre, args.n_per), centres)
np.save("variances.{}c.{}pc.npy".format(args.n_centre, args.n_per), variances)
